import numpy as np
import matplotlib.pyplot as plt

# ----------------------------------------------------------------
# 1. 准备数据：这里只填入了 ALVC (proposed) 的数据，其它算法请按同样方式补全
# ----------------------------------------------------------------
data = {
    'JCT-VC Class B': {
        'ALVC (proposed)': {
            'bpp':  [0.05, 0.10, 0.15, 0.20, 0.25, 0.30],
            'psnr': [32.8, 33.3, 34.0, 34.6, 35.0, 35.4],
        },
        # 'VTM 14.1 (default)':    {'bpp': [...], 'psnr': [...]},
        # 'VTM 14.1 (same GOP)':  {...},  # 依此类推
    },
    'JCT-VC Class C': {
        'ALVC (proposed)': {
            'bpp':  [0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40],
            'psnr': [29.0, 30.0, 31.0, 31.8, 32.4, 33.0, 33.5],
        },
        # 其它算法...
    },
    'JCT-VC Class D': {
        'ALVC (proposed)': {
            'bpp':  [0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40],
            'psnr': [29.0, 30.0, 31.0, 32.0, 32.8, 33.4, 33.8],
        },
    },
    'JCT-VC Class E': {
        'ALVC (proposed)': {
            'bpp':  [0.005, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07],
            'psnr': [35.0, 36.0, 36.8, 37.5, 38.2, 39.0, 40.0],
        },
    },
    'UVG': {
        'ALVC (proposed)': {
            'bpp':  [0.03, 0.05, 0.08, 0.11, 0.14, 0.16],
            'psnr': [34.3, 35.2, 36.0, 36.7, 37.4, 37.8],
        },
    },
    'VTL': {
        'ALVC (proposed)': {
            'bpp':  [0.08, 0.12, 0.16, 0.20, 0.24, 0.28],
            'psnr': [31.5, 32.3, 33.0, 33.6, 34.0, 34.5],
        },
    },
}

# ----------------------------------------------------------------
# 2. 定义画图样式（如需调整可自行改）
# ----------------------------------------------------------------
styles = {
    'ALVC (proposed)':       {'color': 'red',    'linestyle': '-',  'marker': 'o'},
    'VTM 14.1 (default)':     {'color': 'red',    'linestyle': '--','marker': None},
    'VTM 14.1 (same GOP)':   {'color': 'red',    'linestyle': '-.', 'marker': None},
    'HM 16.20 (default)':    {'color': 'purple', 'linestyle': '--','marker': None},
    'HM 16.20 (same GOP)':   {'color': 'purple', 'linestyle': '-.', 'marker': None},
    'x265 (B-frame default)':{'color': 'orange', 'linestyle': '--','marker': None},
    'x265 (LDP default)':    {'color': 'green',  'linestyle': '--','marker': None},
    'x265 (LDP very fast)':  {'color': 'green',  'linestyle': '-.', 'marker': None},
}

# ----------------------------------------------------------------
# 3. 画图
# ----------------------------------------------------------------
fig, axes = plt.subplots(3, 2, figsize=(12, 14))
axes = axes.flatten()

for ax, (title, algos) in zip(axes, data.items()):
    for name, ds in algos.items():
        bpp  = np.array(ds['bpp'])
        psnr = np.array(ds['psnr'])
        ax.plot(bpp, psnr,
                label=name,
                **styles.get(name, {}))
    ax.set_title(title, fontsize=12)
    ax.set_xlabel('bpp', fontsize=10)
    ax.set_ylabel('PSNR (dB)', fontsize=10)
    ax.grid(which='both', linestyle=':', alpha=0.6)
    ax.legend(fontsize=8, loc='lower right')

plt.tight_layout()
plt.savefig("./res.png", format="png", dpi=600)

